import requests
import json
import time
import base64
from io import BytesIO
from PIL import Image

# ComfyUI API 的基础 URL（根据您的实际部署调整）
COMFYUI_API_URL = "http://127.0.0.1:8188"

# 您提供的工作流
workflow = {
    "1": {
        "inputs": {
            "ckpt_name": "animagineXLV31_v30.safetensors"
        },
        "class_type": "CheckpointLoaderSimple",
        "_meta": {
            "title": "Checkpoint加载器(简易)"
        }
    },
    "2": {
        "inputs": {
            "seed": 272052676331338,
            "steps": 20,
            "cfg": 8,
            "sampler_name": "euler",
            "scheduler": "normal",
            "denoise": 1,
            "model": ["1", 0],
            "positive": ["3", 0],
            "negative": ["12", 0],
            "latent_image": ["7", 0]
        },
        "class_type": "KSampler",
        "_meta": {
            "title": "K采样器"
        }
    },
    "3": {
        "inputs": {
            "text": "a girl",
            "clip": ["1", 1]
        },
        "class_type": "CLIPTextEncode",
        "_meta": {
            "title": "正向提示词"
        }
    },
    "7": {
        "inputs": {
            "width": 512,
            "height": 512,
            "batch_size": 1
        },
        "class_type": "EmptyLatentImage",
        "_meta": {
            "title": "空Latent"
        }
    },
    "10": {
        "inputs": {
            "samples": ["2", 0],
            "vae": ["1", 2]
        },
        "class_type": "VAEDecode",
        "_meta": {
            "title": "VAE解码"
        }
    },
    "11": {
        "inputs": {
            "filename_prefix": "ComfyUI",
            "images": ["10", 0]
        },
        "class_type": "SaveImage",
        "_meta": {
            "title": "保存图像"
        }
    },
    "12": {
        "inputs": {
            "text": "text",
            "clip": ["1", 1]
        },
        "class_type": "CLIPTextEncode",
        "_meta": {
            "title": "CLIP文本编码器"
        }
    }
}

def queue_prompt(prompt):
    """将工作流提交到 ComfyUI 的队列"""
    url = f"{COMFYUI_API_URL}/prompt"
    payload = {"prompt": prompt}
    response = requests.post(url, json=payload)
    return response.json()

def get_image(filename, subfolder, folder_type):
    """从 ComfyUI 获取生成的图像"""
    url = f"{COMFYUI_API_URL}/view"
    params = {
        "filename": filename,
        "subfolder": subfolder,
        "type": folder_type
    }
    response = requests.get(url, params=params)
    return response.content

def get_history(prompt_id):
    """获取指定 prompt 的历史记录"""
    url = f"{COMFYUI_API_URL}/history/{prompt_id}"
    response = requests.get(url)
    return response.json()

def main():
    # 提交工作流
    result = queue_prompt(workflow)
    prompt_id = result.get("prompt_id")
    
    if not prompt_id:
        print("提交工作流失败:", result)
        return
    
    print(f"工作流已提交，Prompt ID: {prompt_id}")
    
    # 轮询检查任务状态
    while True:
        history = get_history(prompt_id)
        if prompt_id in history:
            # 任务完成
            break
        print("任务处理中...")
        time.sleep(1)
    
    # 获取输出结果
    outputs = history[prompt_id]["outputs"]
    if "11" in outputs:  # 检查 SaveImage 节点 (ID: 11)
        images = outputs["11"]["images"]
        for image_info in images:
            filename = image_info["filename"]
            subfolder = image_info.get("subfolder", "")
            folder_type = image_info.get("type", "output")
            
            # 获取图像数据
            image_data = get_image(filename, subfolder, folder_type)
            
            # 将图像数据转换为 PIL 图像
            image = Image.open(BytesIO(image_data))
            
            # 保存到本地（可选）
            image.save(f"output_{filename}")
            print(f"图像已保存: output_{filename}")
    
    print("任务完成！")

if __name__ == "__main__":
    main()